import torch
import numpy as np
from loader import load_data
import matplotlib.pyplot as plt
import os
from scipy.stats import pearsonr


class Evaluator:
    def __init__(self, config, model, logger):
        self.config = config
        self.model = model
        self.logger = logger
        self.valid_data = load_data(config["valid_data_path"], config, shuffle=False)

    def eval(self, epoch):
        self.logger.info("开始测试第%d轮模型效果：" % epoch)
        self.score = []
        self.model.eval()
        for index, batch_data in enumerate(self.valid_data):
            if torch.cuda.is_available():
                batch_data = [d.cuda() for d in batch_data]
            input_id, target_id, labels = batch_data
            with torch.no_grad():
                eval_score = self.model(input_id, target_id)  # 将测试数据进行编码，不输入labels，使用模型当前参数进行预测
            temp_score = self.pearsSim(eval_score, labels)
            self.score.append(temp_score)
        self.logger.info("epoch average pearson: %f" % np.mean(self.score))
        return np.mean(self.score)

    def pearsSim(self, tensor1, tensor2):
        array1 = np.array(tensor1.detach().cpu())
        array2 = np.array(tensor2.detach().cpu())
        # pc = 0.5 + 0.5 * np.corrcoef(array1, array2, rowvar=0)
        # pc = pc[0][1]
        array1, array2 = array1.squeeze(), array2.squeeze()
        pc_2 = pearsonr(array1, array2)
        # pc_2 = 0.5 * (1 + pc_2[0])
        pc = pc_2[0]
        # pc = torch.as_tensor(pc, dtype=torch.float32, device='cuda:0')
        return pc


